// Pavel Zolnikov[http://www.codeproject.com/script/profile/whos_who.asp?id=35980], 2002

using System;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Reflection;
using System.Reflection.Emit;
using System.Globalization;


namespace VirusCleaner.Libary
{
	/// <summary>
	/// Provides aids for an Import Adress Table(IAT)-based API interception.
	/// </summary>
	/// <remarks>
	/// Based on the ReplaceIATEntry function by J.Richter
	/// </remarks>
	public class ApiInterceptor : MarshalByRefObject, IDisposable
	{
		protected Delegate	procedure;
		protected String	callerName;
		protected String	targetName;
		protected String	procedureName;
		protected IntPtr	signatureBytes;
		protected IntPtr	oldAddress;

		/// <summary>
		/// Interception starts once object is constructed.
		/// </summary>
		/// <param name="caller">Name of the module that calls API function.</param>
		/// <param name="target">Name of the modile that implements API function.</param>
		/// <param name="procName">Name of the API function to intercept.</param>
		/// <param name="newProc">Delegate to the function that is to be called instead of intercepted.</param>
		/// <remarks>
		/// newProc delegate must contain only 'blittable' parameters (of types bool, int, IntPtr and so on).
		/// </remarks>
		public ApiInterceptor(String caller, String target, String procName, Delegate newProc)
		{
			procedure = newProc;
			callerName = caller;
			targetName = target;
			procedureName = procName;

			SignatureHelper signature = 
				SignatureHelper.GetMethodSigHelper(
				null,
				procedure.Method.CallingConvention ,
				procedure.Method.ReturnType 
				);

			foreach( ParameterInfo param in procedure.Method.GetParameters() )
			{
				signature.AddArgument( param.ParameterType );
			}
						
			byte[] bytes = signature.GetSignature();
			signatureBytes = Marshal.AllocCoTaskMem(bytes.Length);
			for(int i = 0; i< bytes.Length; i++)
			{
				Marshal.WriteByte( signatureBytes, i, bytes[i]);
			}

			IntPtr newAddress = Marshal.GetUnmanagedThunkForManagedMethodPtr(
				procedure.Method.MethodHandle.GetFunctionPointer(),
				signatureBytes,
				bytes.Length
				);

			ReplaceIatEntry(caller, target, procName, newAddress);
		}

		protected virtual void Dispose(bool disposing)
		{
			GC.KeepAlive(this);//FxCop requirement
			ReplaceIatEntry(callerName, targetName, procedureName, oldAddress);
			Marshal.FreeCoTaskMem(signatureBytes);
		}

		/// <summary>
		/// Stops interception, disposes of resources.
		/// </summary>
		public void Dispose()
		{
			Dispose(true);
			GC.SuppressFinalize(this);
		}

		~ApiInterceptor()
		{
			Dispose(false);
		}

		/// <summary>
		/// </summary>
		/// <remarks>
		/// See also J.Richter's "Programming Applications for MS Windows 4th Edition".
		/// </remarks>
		protected void ReplaceIatEntry(String caller, String target, String procName, IntPtr newProc)
		{
			GC.KeepAlive(this);

			Int32 hmodCaller = (Int32)Win32.GetModuleHandle(caller);
			Int32 hmodTarget = (Int32)Win32.GetModuleHandle(target);
			Int32 pfnCurrent = (Int32)Win32.GetProcAddress( (IntPtr)hmodTarget, procName);

			UInt32 Size;
			Int32 pImpDesc = (Int32)ImageDirectoryEntryToData(
				(IntPtr)hmodCaller,
				true, 
				IMAGE_DIRECTORY_ENTRY_IMPORT,
				out Size );

			if( pImpDesc == 0 ) return;  // This module has no import section.


			// Find the import descriptor containing references 
			// to callee's functions.
			IMAGE_IMPORT_DESCRIPTOR ImportDesc = (IMAGE_IMPORT_DESCRIPTOR)Marshal.PtrToStructure((IntPtr)pImpDesc,typeof(IMAGE_IMPORT_DESCRIPTOR) );
			while( ImportDesc.Name != 0 )
			{
				String ModName = Marshal.PtrToStringAnsi((IntPtr)(hmodCaller + ImportDesc.Name));
					
				if( 0 == String.Compare(ModName, target, true, CultureInfo.InvariantCulture) )  break;

				pImpDesc += Marshal.SizeOf(typeof(IMAGE_IMPORT_DESCRIPTOR));
				ImportDesc = (IMAGE_IMPORT_DESCRIPTOR)Marshal.PtrToStructure((IntPtr)pImpDesc,typeof(IMAGE_IMPORT_DESCRIPTOR) );
			}

			if( ImportDesc.Name == 0 ) return; // This module doesn't import any functions from this callee.

			// Get caller's import address table (IAT) 
			// for the callee's functions.
			Int32 pThunk = hmodCaller + ImportDesc.FirstThunk;

			IMAGE_THUNK_DATA32 Thunk = (IMAGE_THUNK_DATA32)Marshal.PtrToStructure((IntPtr)pThunk,typeof(IMAGE_THUNK_DATA32));
			// Replace current function address with new function address.
			while( Thunk.Address!= 0 ) 
			{
				if( Thunk.Address == pfnCurrent )
				{						
					// The addresses match; change the import section address.
					oldAddress = ExchangeIntPtrs((IntPtr)pThunk, newProc);

					return;  // We did it; get out.
				}
					
				pThunk += Marshal.SizeOf(typeof(IMAGE_THUNK_DATA32));
				Thunk = (IMAGE_THUNK_DATA32)Marshal.PtrToStructure((IntPtr)pThunk,typeof(IMAGE_THUNK_DATA32));
			}

			// If we get to here, the function
			// is not in the caller's import section.

		}

		static IntPtr ExchangeIntPtrs(IntPtr valueAddres, IntPtr newValue)
		{
			IntPtr buffer = Marshal.AllocCoTaskMem( Marshal.SizeOf(typeof(IntPtr)) );
			Marshal.WriteIntPtr( buffer, newValue );

			IntPtr old = Marshal.ReadIntPtr(valueAddres);

			Int32 num;
			Win32.WriteProcessMemory(
				Process.GetCurrentProcess().Handle,
				valueAddres,
				buffer,
				Marshal.SizeOf(typeof(IntPtr)),
				out num );

			Marshal.FreeCoTaskMem( buffer );

			return old;
		}


		[DllImport("dbghelp.dll")]
		extern static IntPtr ImageDirectoryEntryToData(
			IntPtr Base,            
			bool MappedAsImage,  
			UInt16 DirectoryEntry,  
			out UInt32 Size            
			);

		const UInt16 IMAGE_DIRECTORY_ENTRY_IMPORT = 1;

		[StructLayout(LayoutKind.Sequential)]
			class IMAGE_IMPORT_DESCRIPTOR
		{
			public Int32   Characteristics_or_OriginalFirstThunk = 0;	// 0 for terminating null import descriptor
			// RVA to original unbound IAT (PIMAGE_THUNK_DATA)
				
			public Int32   TimeDateStamp = 0;	// 0 if not bound,
			// -1 if bound, and real date\time stamp
			//     in IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT (new BIND)
			// O.W. date/time stamp of DLL bound to (Old BIND)

			public Int32   ForwarderChain = 0;	// -1 if no forwarders
			public Int32   Name = 0;
			public Int32   FirstThunk = 0;      // RVA to IAT (if bound this IAT has actual addresses)
		}

		[StructLayout(LayoutKind.Sequential)]
			struct IMAGE_THUNK_DATA32 
		{
			public Int32 Address;//works only on Intel32
		} 
	}

}